"""
计算ATI数据
需要输入地表温度和反照率
"""
import datetime
import pytz
import math
import os.path

import numpy as np
from osgeo import gdal, osr


# 写入tif
def write_tif(file_path, data, geotransform, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    # 相当于一个创建数据集的驱动
    driver = gdal.GetDriverByName("GTiff")
    # 根据数据的维度来确定行列数（图像大小）
    rows, cols = data.shape
    # 创建一个新的数据集，存储输出文件；1是波段
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    # 设置仿射变换矩阵
    dataset.SetGeoTransform(geotransform)
    # 设置投影
    prj = osr.SpatialReference()
    # WGS84投影
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    # 1个波段
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    # 释放内存
    del dataset


# 获取某一年某一天的UTC
def get_utc_by_date(doy, hour, minute):
    year = int(doy[0:4])
    day = int(doy[4:])
    # 将指定年份和天数转为datetime对象
    date = datetime.datetime(year, 1, 1) + datetime.timedelta(day - 1)
    utc = pytz.utc
    # 指定当天的某个时间
    date = date.replace(hour=hour, minute=minute, second=0, microsecond=0)
    # 转为UTC时间
    date_utc = date.replace(tzinfo=utc).timestamp()
    return date_utc


# 计算太阳校正系数
def get_sun_correction_factor(phi, nd):
    # 公式： C = sinφsinδ(1 - tan²φtan²δ)^0.5 + cosφcosδarccos(-tanφtanδ)
    # δ = 0.006918 - 0.399912cos(Γ) + 0.070257sin(Γ) - 0.006758cos(2Γ) + 0.000907sin(2Γ) - 0.002697cos(3Γ) + 0.00148sin(3Γ)
    # Γ = (2 * π * (nd - 1)) / 365.25
    # φ是纬度 δ为太阳赤纬 Γ为昼角 nd为日数
    gamma = (2 * np.pi * (nd - 1)) / 365.25
    delta = 0.006918 - 0.399912 * np.cos(gamma) + 0.070257 * np.sin(gamma) \
            - 0.006758 * np.cos(2 * gamma) + 0.000907 * np.sin(2 * gamma) \
            - 0.002697 * np.cos(3 * gamma) + 0.00148 * np.sin(3 * gamma)
    # 百度百科
    # delta = np.arcsin(0.39795 * np.cos(0.98563 * (nd - 173)))
    print(gamma)
    print(delta * 180 / np.pi)
    C = np.sin(phi) * np.sin(delta) * (1 - np.tan(phi) ** 2 * np.tan(delta) ** 2) ** 0.5 + \
        np.cos(phi) * np.cos(delta) * np.arccos(-np.tan(phi) * np.tan(delta))
    return C


# 计算LST振幅A
def get_lst_amplitude(terra_lst_day, terra_lst_night, aqua_lst_day, aqua_lst_night, omega, doy):
    # 公式：A = 2*((n * sum(cos(ωti - Ψ) * Ti - sum(cos(ωti - Ψ)) * sum(Ti)))/(n * sum(cos²(ωti - Ψ))²))
    # Ti表示时间ti的LST
    # n表示MODIS每天重建的每个像素LST值的个数
    # 相位角
    T1 = terra_lst_day.ReadAsArray()
    T2 = terra_lst_night.ReadAsArray()
    T3 = aqua_lst_day.ReadAsArray()
    T4 = aqua_lst_night.ReadAsArray()

    geo1 = terra_lst_day.GetGeoTransform()
    geo2 = terra_lst_night.GetGeoTransform()
    geo3 = aqua_lst_day.GetGeoTransform()
    geo4 = aqua_lst_night.GetGeoTransform()

    # terra过境时间 10:30AM 10:30PM
    # aqua过境时间 1:30AM   1:30PM
    utc1 = get_utc_by_date(doy, 10, 30)
    utc2 = get_utc_by_date(doy, 22, 30)
    utc3 = get_utc_by_date(doy, 1, 30)
    utc4 = get_utc_by_date(doy, 13, 30)
    # print(utc1)
    # print(utc2)
    # print(utc3)
    # print(utc4)

    lon_arr1 = np.zeros_like(T1, dtype=np.float32)
    lon_arr2 = np.zeros_like(T2, dtype=np.float32)
    lon_arr3 = np.zeros_like(T3, dtype=np.float32)
    lon_arr4 = np.zeros_like(T4, dtype=np.float32)
    for i in range(T1.shape[0]):
        for j in range(T1.shape[1]):
            lon_arr1[i, j] = math.radians(geo1[0] + geo1[1] * i + geo1[2] * j)
            lon_arr2[i, j] = math.radians(geo2[0] + geo2[1] * i + geo2[2] * j)
            lon_arr3[i, j] = math.radians(geo3[0] + geo3[1] * i + geo3[2] * j)
            lon_arr4[i, j] = math.radians(geo4[0] + geo4[1] * i + geo4[2] * j)

    t1 = utc1 + lon_arr1 / omega
    t2 = utc2 + lon_arr2 / omega
    t3 = utc3 + lon_arr3 / omega
    t4 = utc4 + lon_arr4 / omega
    # print(t1)
    n = 4
    t = [t1, t2, t3, t4]
    T = [T1, T2, T3, T4]

    # 计算相位角
    a = (T1 - T3) * (np.cos(omega * t2) - np.cos(omega * t4)) - (T2 - T4) * (np.cos(omega * t1) - np.cos(omega * t3))
    b = (T2 - T4) * (np.sin(omega * t1) - np.sin(omega * t3)) - (T1 - T3) * (np.cos(omega * t2) - np.cos(omega * t4))
    print(a)
    print(b)
    xi = a / b
    puxi = np.arctan(xi) + np.pi

    # 计算振幅A
    c = 0
    d = 0
    e = 0
    f = 0
    g = 0
    for i in range(n):
        c = c + np.cos(omega * t[i] - puxi) * T[i]
        d = d + np.cos(omega * t[i] - puxi)
        g = g + T[i]
        e = e + np.cos(omega * t[i] - puxi) ** 2
        f = f + np.cos(omega * t[i] - puxi)
    # c = n * np.sum(np.cos(omega * t - puxi) * T) - np.sum(np.cos(omega * t - puxi) * np.sum(T))
    # d = n * np.sum(np.cos(omega * t - puxi) ** 2) - np.sum(math.cos(omega * t - puxi)) ** 2
    A = 2 * (n * c - d * g) / (n * e - f ** 2)
    return A


# 计算ATI
# 有两个计算公式
# 1. ATI=(1-ABE)/ΔT
# 2. ATI=C * (1 - a) / A  C是太阳校正因子，a是实际地表反照率，A是昼夜LST周期振幅
def get_ati(terra_lst_day_path, terra_lst_night_path, aqua_lst_day_path, aqua_lst_night_path, albedo_bsa_path, albedo_wsa_path, snow_path, water_path, output_path):
    terra_lst_day = gdal.Open(terra_lst_day_path)
    terra_lst_night = gdal.Open(terra_lst_night_path)
    aqua_lst_day = gdal.Open(aqua_lst_day_path)
    aqua_lst_night = gdal.Open(aqua_lst_night_path)
    albedo_bsa = gdal.Open(albedo_bsa_path)
    albedo_wsa = gdal.Open(albedo_wsa_path)

    terra_lst_day_array = terra_lst_day.ReadAsArray()
    terra_lst_night_array = terra_lst_night.ReadAsArray()
    aqua_lst_day_array = aqua_lst_day.ReadAsArray()
    aqua_lst_night_array = aqua_lst_night.ReadAsArray()
    albedo_bsa_array = albedo_bsa.ReadAsArray()
    albedo_wsa_array = albedo_wsa.ReadAsArray()
    snow_array = gdal.Open(snow_path).ReadAsArray(xsize=terra_lst_day_array.shape[0], ysize=terra_lst_day_array.shape[1])
    water_array = gdal.Open(water_path).ReadAsArray(xsize=terra_lst_day_array.shape[0], ysize=terra_lst_day_array.shape[1])
    # 雪和水体掩膜
    mask = (snow_array != 0) & (water_array != 1)
    terra_lst_day_array = np.where(mask, np.nan, terra_lst_day_array)
    terra_lst_night_array = np.where(mask, np.nan, terra_lst_night_array)
    aqua_lst_day_array = np.where(mask, np.nan, aqua_lst_day_array)
    aqua_lst_night_array = np.where(mask, np.nan, aqua_lst_night_array)
    albedo_bsa_array = np.where(mask, np.nan, albedo_bsa_array)
    albedo_wsa_array = np.where(mask, np.nan, albedo_wsa_array)

    albedo_bsa_array = np.where(-1, np.nan, albedo_bsa_array)
    albedo_wsa_array = np.where(-1, np.nan, albedo_wsa_array)
    albedo_mean = np.nanmean([albedo_bsa_array, albedo_wsa_array], axis=0) * 0.0001

    # albedo_mean = (albedo_bsa_array + albedo_wsa_array) * 0.0001 / 2

    # print(albedo_mean)
    geotrans = terra_lst_day.GetGeoTransform()

    doy = terra_lst_day_path.split("\\")[-1].split(".")[0]
    # print(doy)
    # 计算地表温度振幅
    # omega = 7.292e-5
    # A = get_lst_amplitude(terra_lst_day, terra_lst_night, aqua_lst_day, aqua_lst_night, omega, doy)
    # print(A.mean())

    # 计算太阳校正系数
    nd = int(doy[-3:])
    # print(nd)
    # print(geotrans)
    lat_arr = np.zeros_like(aqua_lst_day_array, dtype=np.float32)
    for i in range(aqua_lst_day_array.shape[0]):
        for j in range(aqua_lst_day_array.shape[1]):
            lat_arr[i, j] = math.radians(geotrans[3] + geotrans[4] * i + geotrans[5] * j)
    # print(lat_arr)
    C = get_sun_correction_factor(lat_arr, nd)

    # print(C)
    # 计算昼夜温度差
    print(np.nanmax(aqua_lst_day_array))
    print(np.nanmin(aqua_lst_night_array[aqua_lst_night_array != 0]))
    # delta_t = (np.nanmax(aqua_lst_day_array) - np.nanmin(aqua_lst_night_array[aqua_lst_night_array != 0])) * 0.02
    delta_t = (np.nanmax(terra_lst_day_array) - np.nanmin(terra_lst_night_array[terra_lst_night_array != 0]))
    # print(delta_t)
    print("太阳校正系数：", C)  # 1.5
    print("昼夜温差：", delta_t)  # 50
    # print("平均反照率：", albedo_mean)    # 0.1
    ati_array = C * (1 - albedo_mean) / delta_t     # 0-0.79
    write_tif(output_path, ati_array, geotrans, nodata=-9999, dataType='float32')


def main(terra_lst_day_dir, terra_lst_night_dir, aqua_lst_day_dir, aqua_lst_night_dir, albedo_bsa_dir, albedo_wsa_dir, snow_dir, water_dir, output_dir):
    # lst_day_list = [os.path.join(lst_day_dir, x) for x in os.listdir(lst_day_dir)]
    # lst_night_list = [os.path.join(lst_night_dir, x) for x in os.listdir(lst_night_dir)]
    # albedo_bsa_list = [os.path.join(albedo_bsa_dir, x) for x in os.listdir(albedo_bsa_dir)]
    # albedo_wsa_list = [os.path.join(albedo_wsa_dir, x) for x in os.listdir(albedo_wsa_dir)]
    lst_file_list = os.listdir(terra_lst_day_dir)
    for file in lst_file_list:
        terra_lst_day_path = os.path.join(terra_lst_day_dir, file)
        terra_lst_night_path = os.path.join(terra_lst_night_dir, file)
        aqua_lst_day_path = os.path.join(aqua_lst_day_dir, file)
        aqua_lst_night_path = os.path.join(aqua_lst_night_dir, file)
        albedo_bsa_path = os.path.join(albedo_bsa_dir, file)
        albedo_wsa_path = os.path.join(albedo_wsa_dir, file)
        snow_path = os.path.join(snow_dir, file)
        water_path = os.path.join(water_dir, file)
        output_path = os.path.join(output_dir, file)
        get_ati(terra_lst_day_path, terra_lst_night_path, aqua_lst_day_path, aqua_lst_night_path, albedo_bsa_path, albedo_wsa_path, snow_path, water_path, output_path)
    print("计算完成！")


if __name__ == '__main__':
    terra_lst_day_dir = r'G:\test\process_result\LST_rec_tif\terra\clip\day'
    terra_lst_night_dir = r'G:\test\process_result\LST_rec_tif\terra\clip\night'
    aqua_lst_day_dir = r'G:\test\process_result\LST_rec_tif\aqua\clip\day'
    aqua_lst_night_dir = r'G:\test\process_result\LST_rec_tif\aqua\clip\night'

    # terra_lst_day_dir = r'G:\test\process_result\LST_tif\terra\merge_clip\LST_day'
    # terra_lst_night_dir = r'G:\test\process_result\LST_tif\terra\merge_clip\LST_night'
    # aqua_lst_day_dir = r'G:\test\process_result\LST_tif\aqua\merge_clip\LST_day'
    # aqua_lst_night_dir = r'G:\test\process_result\LST_tif\aqua\merge_clip\LST_night'

    albedo_bsa_dir = r"G:\test\process_result\GLASS_tif\glass_linear\ABD_BSA_VIS"
    albedo_wsa_dir = r"G:\test\process_result\GLASS_tif\glass_linear\ABD_WSA_VIS"
    snow_dir = r"G:\test\process_result\MCD_tif\snow"
    water_dir = r"G:\test\process_result\MCD_tif\water"
    output_dir = r"G:\test\process_result\ATI_tif"
    main(terra_lst_day_dir, terra_lst_night_dir, aqua_lst_day_dir, aqua_lst_night_dir, albedo_bsa_dir, albedo_wsa_dir, snow_dir, water_dir, output_dir)
